#!/bin/bash

# Qwen2.5-32B 环境测试脚本
# 用于验证部署是否成功

set -e

echo "========================================="
echo "开始测试 Qwen2.5-32B 环境..."
echo "========================================="

# 颜色定义
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 检查计数器
PASSED=0
FAILED=0

# 检查函数
check_item() {
    local name=$1
    local check_cmd=$2
    
    echo -n "检查 $name... "
    if eval "$check_cmd" > /dev/null 2>&1; then
        echo -e "${GREEN}✓ 通过${NC}"
        ((PASSED++))
        return 0
    else
        echo -e "${RED}✗ 失败${NC}"
        ((FAILED++))
        return 1
    fi
}

# 1. 检查 conda 是否可用
echo ""
echo "【步骤 1】检查 Conda 环境..."
if ! command -v conda &> /dev/null; then
    if [ -f "$HOME/miniconda/bin/conda" ]; then
        export PATH="$HOME/miniconda/bin:$PATH"
        echo "已加载 Miniconda 路径"
    else
        echo -e "${RED}错误: 未找到 conda 命令${NC}"
        echo "请先运行: source ~/.bashrc 或 export PATH=\"\$HOME/miniconda/bin:\$PATH\""
        exit 1
    fi
fi

check_item "Conda 命令" "conda --version"

# 检查 conda 是否已初始化
if ! conda info --envs &> /dev/null && ! command -v conda &> /dev/null; then
    echo -e "${YELLOW}⚠ 检测到 Conda 可能未初始化${NC}"
    echo "正在尝试初始化 Conda..."
    if [ -n "$BASH_VERSION" ]; then
        conda init bash > /dev/null 2>&1
        echo "已初始化 bash，请运行: source ~/.bashrc"
    elif [ -n "$ZSH_VERSION" ]; then
        conda init zsh > /dev/null 2>&1
        echo "已初始化 zsh，请运行: source ~/.zshrc"
    fi
fi

# 2. 检查 qwen32b 环境是否存在
echo ""
echo "【步骤 2】检查 Conda 虚拟环境..."
# 尝试使用 conda env list，如果失败则尝试其他方法
if conda env list 2>/dev/null | grep -q "qwen32b"; then
    echo -e "${GREEN}✓ qwen32b 环境存在${NC}"
    ((PASSED++))
elif [ -d "$HOME/miniconda/envs/qwen32b" ] || [ -d "$HOME/anaconda3/envs/qwen32b" ] || [ -d "$CONDA_PREFIX/../envs/qwen32b" ]; then
    echo -e "${GREEN}✓ qwen32b 环境存在（通过路径检测）${NC}"
    ((PASSED++))
else
    echo -e "${RED}✗ qwen32b 环境不存在${NC}"
    echo "请先运行 deploy_qwen.sh 脚本"
    exit 1
fi

# 3. 激活环境并检查 Python 版本
echo ""
echo "【步骤 3】激活环境并检查 Python..."
# 尝试初始化 conda（如果还没有）
if ! type conda | grep -q function; then
    # 如果 conda 不是函数，尝试初始化
    if [ -f "$HOME/miniconda/etc/profile.d/conda.sh" ]; then
        source "$HOME/miniconda/etc/profile.d/conda.sh"
    elif [ -f "$HOME/anaconda3/etc/profile.d/conda.sh" ]; then
        source "$HOME/anaconda3/etc/profile.d/conda.sh"
    else
        # 尝试使用 conda init
        eval "$(conda shell.bash hook 2>/dev/null)" || {
            echo -e "${YELLOW}⚠ 无法自动初始化 conda${NC}"
            echo "请手动运行以下命令之一："
            echo "  1. source ~/.bashrc"
            echo "  2. conda init bash && source ~/.bashrc"
            echo "  3. 或者使用: $HOME/miniconda/bin/conda activate qwen32b"
            # 尝试直接使用 conda 可执行文件激活
            if [ -f "$HOME/miniconda/bin/conda" ]; then
                echo "尝试直接使用 conda 可执行文件..."
                export PATH="$HOME/miniconda/bin:$PATH"
                CONDA_BASE="$HOME/miniconda"
                source "$CONDA_BASE/etc/profile.d/conda.sh"
            else
                exit 1
            fi
        }
    fi
fi

# 激活环境
if conda activate qwen32b 2>/dev/null; then
    echo -e "${GREEN}✓ 环境激活成功${NC}"
elif [ -n "$CONDA_PREFIX" ]; then
    # 如果已经在某个 conda 环境中，尝试切换到 qwen32b
    conda deactivate 2>/dev/null || true
    conda activate qwen32b 2>/dev/null || {
        echo -e "${YELLOW}⚠ 使用备用方法激活环境${NC}"
        # 备用方法：直接设置环境变量
        if [ -d "$HOME/miniconda/envs/qwen32b" ]; then
            export CONDA_DEFAULT_ENV=qwen32b
            export CONDA_PREFIX="$HOME/miniconda/envs/qwen32b"
            export PATH="$CONDA_PREFIX/bin:$PATH"
            echo -e "${GREEN}✓ 环境已通过备用方法激活${NC}"
        else
            echo -e "${RED}✗ 无法激活环境${NC}"
            exit 1
        fi
    }
else
    echo -e "${RED}✗ 无法激活环境${NC}"
    echo ""
    echo "建议运行修复脚本："
    echo "  bash $(dirname "$0")/fix_conda_init.sh"
    echo ""
    echo "或手动运行："
    echo "  source ~/.bashrc"
    echo "  conda activate qwen32b"
    exit 1
fi

check_item "Python 版本 (应为 3.10)" "python --version | grep 'Python 3.10'"

# 4. 检查必要的 Python 包
echo ""
echo "【步骤 4】检查 Python 依赖包..."
check_item "PyTorch" "python -c 'import torch; print(torch.__version__)'"
check_item "vLLM" "python -c 'import vllm; print(vllm.__version__)'"
check_item "ModelScope" "python -c 'import modelscope'"

# 5. 检查 CUDA 是否可用
echo ""
echo "【步骤 5】检查 CUDA 支持..."
if python -c "import torch; print('CUDA available:', torch.cuda.is_available())" 2>/dev/null | grep -q "CUDA available: True"; then
    echo -e "${GREEN}✓ CUDA 可用${NC}"
    CUDA_VERSION=$(python -c "import torch; print(torch.version.cuda)" 2>/dev/null)
    GPU_COUNT=$(python -c "import torch; print(torch.cuda.device_count())" 2>/dev/null)
    echo "  CUDA 版本: $CUDA_VERSION"
    echo "  GPU 数量: $GPU_COUNT"
    ((PASSED++))
else
    echo -e "${YELLOW}⚠ CUDA 不可用（如果使用 CPU 推理可忽略）${NC}"
    ((FAILED++))
fi

# 6. 检查模型文件是否存在
echo ""
echo "【步骤 6】检查模型文件..."
MODEL_DIR="./models/Qwen/Qwen2.5-32B"
if [ -d "$MODEL_DIR" ]; then
    echo -e "${GREEN}✓ 模型目录存在: $MODEL_DIR${NC}"
    MODEL_SIZE=$(du -sh "$MODEL_DIR" 2>/dev/null | cut -f1)
    echo "  模型大小: $MODEL_SIZE"
    ((PASSED++))
    
    # 检查关键文件
    if [ -f "$MODEL_DIR/config.json" ]; then
        echo -e "${GREEN}✓ 模型配置文件存在${NC}"
        ((PASSED++))
    else
        echo -e "${YELLOW}⚠ 模型配置文件不存在${NC}"
        ((FAILED++))
    fi
else
    echo -e "${RED}✗ 模型目录不存在: $MODEL_DIR${NC}"
    echo "  请检查模型是否已下载完成"
    ((FAILED++))
fi

# 7. 快速推理测试（可选，需要 GPU）
echo ""
echo "【步骤 7】快速推理测试..."
echo "注意: 此测试需要 GPU 和足够的内存，可能需要几分钟时间"
read -p "是否执行推理测试? (y/N): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    echo "正在测试模型加载和推理..."
    
    TEST_SCRIPT=$(cat <<'EOF'
import sys
import torch

# 首先检查 CUDA
if not torch.cuda.is_available():
    print("✗ 错误: PyTorch 未检测到 CUDA")
    print("请先检查 GPU 环境，可以运行诊断脚本:")
    print("  bash $(dirname "$0")/diagnose_gpu.sh")
    sys.exit(1)

try:
    from vllm import LLM, SamplingParams
    import os
    
    model_path = "./models/Qwen/Qwen2.5-32B"
    if not os.path.exists(model_path):
        print("错误: 模型路径不存在")
        sys.exit(1)
    
    print("正在加载模型（这可能需要几分钟）...")
    llm = LLM(model=model_path, trust_remote_code=True, max_model_len=2048)
    
    print("模型加载成功！")
    print("正在执行测试推理...")
    
    prompts = ["你好，请介绍一下你自己。"]
    sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=50)
    outputs = llm.generate(prompts, sampling_params)
    
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"\n提示: {prompt}")
        print(f"回复: {generated_text}")
    
    print("\n✓ 推理测试成功！")
except RuntimeError as e:
    error_msg = str(e)
    if "Device string must not be empty" in error_msg or "No platform detected" in error_msg:
        print("✗ 错误: vLLM 无法检测到 GPU 设备")
        print("这通常意味着 vLLM 无法检测到 CUDA 设备")
        print("\n建议:")
        print("  1. 运行诊断脚本: bash $(dirname \"$0\")/diagnose_gpu.sh")
        print("  2. 检查 nvidia-smi 是否正常工作")
        print("  3. 检查 CUDA 环境变量是否正确设置")
    else:
        print(f"✗ 推理测试失败: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)
except Exception as e:
    print(f"✗ 推理测试失败: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)
EOF
)
    
    if python -c "$TEST_SCRIPT"; then
        echo -e "${GREEN}✓ 推理测试通过${NC}"
        ((PASSED++))
    else
        echo -e "${RED}✗ 推理测试失败${NC}"
        ((FAILED++))
    fi
else
    echo "跳过推理测试"
fi

# 总结
echo ""
echo "========================================="
echo "测试总结"
echo "========================================="
echo -e "${GREEN}通过: $PASSED${NC}"
if [ $FAILED -gt 0 ]; then
    echo -e "${RED}失败: $FAILED${NC}"
else
    echo -e "${GREEN}失败: $FAILED${NC}"
fi

if [ $FAILED -eq 0 ]; then
    echo ""
    echo -e "${GREEN}✓ 所有测试通过！Qwen2.5-32B 环境已准备就绪。${NC}"
    echo ""
    echo "使用示例："
    echo "  1. 激活环境: conda activate qwen32b"
    echo "  2. 启动 vLLM 服务:"
    echo "     python -m vllm.entrypoints.openai.api_server \\"
    echo "       --model ./models/Qwen/Qwen2.5-32B \\"
    echo "       --trust-remote-code"
    echo "  3. 或使用 Python API:"
    echo "     from vllm import LLM, SamplingParams"
    echo "     llm = LLM(model='./models/Qwen/Qwen2.5-32B', trust_remote_code=True)"
    exit 0
else
    echo ""
    echo -e "${YELLOW}⚠ 部分测试失败，请检查上述错误信息${NC}"
    echo ""
    echo "如果遇到 GPU 检测问题，可以运行诊断脚本："
    echo "  bash $(dirname "$0")/diagnose_gpu.sh"
    exit 1
fi

